from zhipuai import ZhipuAI
import json

def init_glm_ai(api_key):
    #api_key = input("智谱AI平台的API_KEY：")
    return  ZhipuAI(api_key=api_key) # 请填写您自己的APIKey

def chatTools(client,prompt,tools,temperature = 0.6):
    response = client.chat.completions.create(
        model="glm-4", # 填写需要调用的模型名称
        messages = [
            {
                "role": "user",
                "content": prompt
            },

        ],
        tools = tools,
        tool_choice="auto",
        temperature=temperature
    )
    return response.choices[0]

def chat_glmOnce(client,prompt):
    response = client.chat.completions.create(
        model="glm-4", # 填写需要调用的模型名称
        messages = [
            {
                "role": "user",
                "content": prompt
            }
        ]
    )
    return response.choices[0].message.content


def ask_glm(client, tools, temperature, content, function_list, dataset):
    ack = chatTools(client, content, tools, temperature)

    if ack.finish_reason == "tool_calls":
        tool_call = ack.message.tool_calls[0]
        function_name = tool_call.function.name
        function_args_dict = json.loads(tool_call.function.arguments)

        # 动态注入参数
        func = function_list[function_name]
        required_params = func.__code__.co_varnames[:func.__code__.co_argcount]

        # 仅注入函数需要的参数
        final_args = {
            k: v for k, v in function_args_dict.items()
            if k in required_params
        }
        final_args["dataset"] = dataset  # 注入数据集

        return func(**final_args)

    return ack.message
